import torch
import torch.nn as nn
import torchvision
from PIL import Image


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.max_pool1 = nn.MaxPool2d(kernel_size=2)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.max_pool2 = nn.MaxPool2d(kernel_size=2)
        self.relu2 = nn.ReLU()
        self.flatten = nn.Flatten()
        self.l1 = nn.Linear(4096, 128)
        self.l2 = nn.Linear(128, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.max_pool2(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.max_pool2(x)
        x = self.relu2(x)
        x = self.flatten(x)
        x = self.l1(x)
        x = self.l2(x)
        x = self.softmax(x)
        return x


model = Model()
model.load_state_dict(torch.load("./model_CIFAR10.pth"))
model.eval()

img = Image.open("../dog2.jpg")
img = img.resize((32, 32))
transform = torchvision.transforms.ToTensor()
img = transform(img)
img = torch.unsqueeze(img, 0)

# print(img.shape)

with torch.no_grad():
    outputs = model.forward(img)
    res = torch.argmax(outputs, dim=1).item()

if res == 0:
    print("airplane")
elif res == 1:
    print("automobile")
elif res == 2:
    print("bird")
elif res == 3:
    print("cat")
elif res == 4:
    print("deer")
elif res == 5:
    print("dog")
elif res == 6:
    print("frog")
elif res == 7:
    print("horse")
elif res == 8:
    print("ship")
elif res == 9:
    print("truck")
else:
    print("oh-no!")
